package actors

import (
	"xagent/internal/messages"
	"xagent/internal/types"
	proto "xagent/proto/v1"

	"github.com/asynkron/protoactor-go/actor"
	"github.com/sirupsen/logrus"
)

// TaskAgent represents an agent specialized in task execution
type TaskAgent struct {
	*BaseAgent
	tasks         map[string]*proto.Task
	taskProcessor types.TaskProcessor
	supervisor    *actor.PID
}

// NewTaskAgent creates a new task agent
func NewTaskAgent(name string, taskProcessor types.TaskProcessor, supervisor *actor.PID) *TaskAgent {
	return &TaskAgent{
		BaseAgent:     NewBaseAgent(name, "Task execution and management specialist"),
		tasks:         make(map[string]*proto.Task),
		taskProcessor: taskProcessor,
		supervisor:    supervisor,
	}
}

// ProcessTask processes a task
func (a *TaskAgent) ProcessTask(protoTask *proto.Task) error {
	// Process task using the processor
	if err := a.taskProcessor.ProcessTask(protoTask); err != nil {
		a.logger.WithError(err).Error("Failed to process task")

		// Send failure message
		message := &proto.AgentMessage{
			Id:   protoTask.Id,
			Type: proto.MessageType_MESSAGE_TYPE_ERROR,
			Content: &proto.AgentMessage_Error{
				Error: &proto.Error{
					Code:    "TASK_PROCESSING_FAILED",
					Message: "Task processing failed",
					Details: map[string]string{
						"error": err.Error(),
					},
				},
			},
			Sender: a.GetRole(),
		}

		if parent := a.GetContext().Parent(); parent != nil {
			a.SendAgentMessage(parent, message)
		}
		return err
	}

	// Store task in local state
	a.tasks[protoTask.Id] = protoTask

	// Send success message
	message := &proto.AgentMessage{
		Id:   protoTask.Id,
		Type: proto.MessageType_MESSAGE_TYPE_RESPONSE,
		Content: &proto.AgentMessage_Message{
			Message: "Task processed successfully",
		},
		Sender: a.GetRole(),
		Metadata: map[string]string{
			"status": "completed",
		},
	}

	if parent := a.GetContext().Parent(); parent != nil {
		a.SendAgentMessage(parent, message)
	}
	return nil
}

// GetCapabilities overrides BaseAgent's GetCapabilities to add task processing capabilities
func (a *TaskAgent) GetCapabilities() []string {
	return append(a.BaseAgent.GetCapabilities(), "task_processing", "resource_management")
}

// processTask handles task processing for the task agent
func (a *TaskAgent) processTask(msg interface{}) {
	switch t := msg.(type) {
	case *proto.Task:
		a.handleTaskExecution(t)
	case *proto.AgentMessage:
		if t.Type == proto.MessageType_MESSAGE_TYPE_REQUEST {
			// Handle task-related message
			a.logger.WithFields(logrus.Fields{
				"task_id": t.Id,
				"content": t.Content,
			}).Info("Processing task message")
		}
	}
}

// handleTaskExecution processes a specific task
func (a *TaskAgent) handleTaskExecution(task *proto.Task) {
	a.logger.WithFields(logrus.Fields{
		"task_id":   task.Id,
		"task_name": task.Name,
	}).Info("Executing task")

	// Store task in local state
	a.tasks[task.Id] = task

	// Implement task execution logic here
	// For now, we'll just mark it as completed
	task.Status = proto.TaskStatus_TASK_STATUS_COMPLETED

	// Notify completion
	a.notifyTaskCompletion(task)
}

// notifyTaskCompletion notifies about task completion
func (a *TaskAgent) notifyTaskCompletion(task *proto.Task) {
	message := &proto.AgentMessage{
		Id:   task.Id,
		Type: proto.MessageType_MESSAGE_TYPE_RESPONSE,
		Content: &proto.AgentMessage_Message{
			Message: "Task completed successfully",
		},
		Sender: a.GetRole(),
		Metadata: map[string]string{
			"status": "completed",
		},
	}

	if parent := a.GetContext().Parent(); parent != nil {
		a.SendAgentMessage(parent, message)
	}
}

// convertToMessage converts AgentMessage to Message interface
func (a *TaskAgent) convertToMessage(msg *proto.AgentMessage) messages.Message {
	var content string
	switch c := msg.Content.(type) {
	case *proto.AgentMessage_Message:
		content = c.Message
	case *proto.AgentMessage_Error:
		content = c.Error.Message
	case *proto.AgentMessage_Prediction:
		content = c.Prediction.String()
	default:
		content = ""
	}

	metadata := make(map[string]interface{})
	for k, v := range msg.Metadata {
		metadata[k] = v
	}

	basicMsg := &messages.BasicMessage{
		MessageID:   msg.Id,
		MsgType:     messages.MessageType(msg.Type.String()),
		MsgContent:  content,
		MsgMetadata: metadata,
	}
	return basicMsg
}

// SendAgentMessage sends a proto.AgentMessage directly
func (a *TaskAgent) SendAgentMessage(target *actor.PID, msg *proto.AgentMessage) {
	a.GetContext().Send(target, msg)
}
